import numpy as np
import PIL.Image as image
from sklearn.cluster import KMeans


def loadData(filePath):
    # 以二进制形式打开文件
    f = open(filePath, 'rb')
    data = []
    img = image.open(f)
    m, n = img.size
    # RGB归一
    for i in range(m):
        for j in range(n):
            x, y, z = img.getpixel((i, j))
            data.append([x / 256.0, y / 256.0, z / 256.0])
    f.close()
    # 以矩阵形式返回
    return np.mat(data), m, n


imgData, row, col = loadData('kmeans/bull.jpg')
# n_clusters 聚类中心个数
# 获得每个像素所属类别
label = KMeans(n_clusters=4).fit_predict(imgData)
# 创建灰度图 保存聚类结果
label = label.reshape([row, col])
pic_new = image.new("L", (row, col))
for i in range(row):
    for j in range(col):
        pic_new.putpixel((i, j), int(256 / (label[i][j] + 1)))
pic_new.save("result-bull-4.jpg", "JPEG")
